import pandas as pd
import datetime
from libs.Objects import Objects
from datetime import timedelta
from StockTradingData import StockTradingData
from libs.StockBase import StockBase
from queue import Queue

DAY_FORMAT = '%Y-%m-%d'


class StockDownloader(object):
    def __init__(self):
        pass

    def getTradingDaysOff(self, days):
        if days > 0:
            raise Exception("days must be <=0")
        today = datetime.date.today()
        endStr = today.strftime(DAY_FORMAT)
        startDay = today + timedelta(days=days * 2)
        startStr = startDay.strftime(DAY_FORMAT)

        return StockBase().getTradingDays(startDay=startStr, endDay=endStr)

    def downloadTradingData(self, startDay, endDay):
        dateList = StockBase().getTradingDays(startDay=startDay, endDay=endDay)
        for date in dateList:
            self.downTradingData(tradingDay=date)

    def downTradingData(self, tradingDay):
        all_files = Objects.getAllFile(dir_path='data/trading')
        if tradingDay in [str_1.replace('data/trading/', '').replace('.csv', '') for str_1 in all_files]:
            print('file %s.csv exists!' % tradingDay)
            return

        data = self.getAvailStock()
        queue = Queue()
        # mutilProcess
        for code in data['code']:
            StockTradingData(tradingDay, code, queue).download()

        print('finished ' + str(queue.qsize()))
        df_arry = [queue.get() for i in range(queue.qsize())]
        all_df = pd.concat(df_arry)
        Objects.mkdir("data/trading")
        all_df.to_csv("data/trading/{date}.csv".format(date=tradingDay), encoding="utf8", index=False)

    ######
    #
    # type	证券类型，其中1：股票，2：指数，3：其它，4：可转债，5：ETF,
    # status	上市状态，其中1：上市，0：退市
    #
    ####
    def downloadStockBasic(self):
        result = StockBase().getStockBasic()
        Objects.mkdir("data")
        result.to_csv("data/stock_basic.csv", encoding="utf8", index=False)

    def getAvailStock(self):
        df = StockBase().getStockBasic()
        lastDay = Objects.getDay(-365)
        result = df[(df['status'] == '1')
                    & (df['type'] == '1')
                    & (~df['code_name'].str.contains('ST'))
                    & (df['ipoDate'] <= lastDay)]
        return result
        # result.to_csv("data/stock_basic_min.csv", encoding="utf8", index=False)

    ### 去掉所有ST股票
    def getAllStock(self, day):
        bs = StockBase.getBaostock()
        rs = bs.query_all_stock(day=day)
        data_list = []
        while (rs.error_code == '0') & rs.next():
            # ['sz.399998', '1', '中证煤炭指数']
            row = rs.get_row_data()
            if row[2].startswith('ST') or row[2].startswith('*ST') or row[2].startswith('N') or row[2].startswith('C'):
                continue
            data_list.append(row[0])
        return data_list


if __name__ == '__main__':
    handler = StockDownloader()
    # 23 web ehhandler.downloadStockBasic()
    # handler.downloadTradingData('2025-01-09', '2025-01-09')
    START_DAY = '2025-04-07'
    END_DAY = '2025-04-08'
    for i in range(100):
        curday = Objects.getDiffDay(START_DAY, i)
        if curday == END_DAY:
            exit()
        handler.downloadTradingData(startDay=curday, endDay=curday)
